//  Copyright (c) Microsoft Corporation.  All Rights Reserved.

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Text;
using System.Xml;
using System.Security.Cryptography.X509Certificates;

using System.ServiceModel.Security;
using System.ServiceModel.Description;
using System.IdentityModel.Selectors;
using System.IdentityModel.Claims;
using System.IdentityModel.Policy;
using System.IdentityModel.Tokens;

using EAI.SAML;

namespace EAI.Common
{
    public class CustomTokenSerializer : WSSecurityTokenSerializer
    {
        public const string WsSecurityPrefix = "wsse";
        public const string WsSecurityNamespace = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd";
        public const string Wsse11Prefix = "wsse11";
        public const string Wsse11SecurityNamespace = "http://docs.oasis-open.org/wss/oasis-wss-wssecurity-secext-1.1.xsd";
        public const string SecurityTokenReference = "SecurityTokenReference";
        public const string TokenType = "TokenType";
        public const string WsUtilityPrefix = "wsu";
        public const string WsUtilityNamespace = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";
        public const string KeyIdentifier = "KeyIdentifier";
        public const string ValueType = "ValueType";

        public const string EAISamlValueType = "http://docs.oasis-open.org/wss/oasis-wss-saml-token-profile-1.1#SAMLID";
        public const string EAISamlTokenType = "http://docs.oasis-open.org/wss/oasis-wss-saml-token-profile-1.1#SAMLV2.0";

        EAISamlSerializer samlSerializer = null;

        public CustomTokenSerializer(SecurityVersion version, EAISamlSerializer samlSerializer)
            : base(version)
        {
            this.samlSerializer = samlSerializer;
        }

        protected override bool CanReadTokenCore(XmlReader reader)
        {
            return (reader.IsStartElement(EAISamlConstants.Assertion, EAISamlConstants.Namespace) || reader.IsStartElement(EAISamlConstants.EncryptedAssertion, EAISamlConstants.Namespace) || base.CanReadTokenCore(reader));
        }

        protected override SecurityToken ReadTokenCore(XmlReader reader, SecurityTokenResolver tokenResolver)
        {
            if (reader.IsStartElement(EAISamlConstants.Assertion, EAISamlConstants.Namespace) ||
                reader.IsStartElement(EAISamlConstants.EncryptedAssertion, EAISamlConstants.Namespace))
            {
                return this.samlSerializer.ReadToken(reader, WSSecurityTokenSerializer.DefaultInstance, tokenResolver);
            }
            else
            {
                return base.ReadTokenCore(reader, tokenResolver);
            }
        }

        protected override bool CanWriteTokenCore(SecurityToken token)
        {
            return ((token is EAISamlSecurityToken) || base.CanWriteTokenCore(token));
        }

        protected override void WriteTokenCore(XmlWriter writer, SecurityToken token)
        {
            if (token is EAISamlSecurityToken)
            {
                samlSerializer.WriteToken((EAISamlSecurityToken)token, writer, WSSecurityTokenSerializer.DefaultInstance);
            }
            else
            {
                base.WriteTokenCore(writer, token);
            }
        }

        protected override bool CanReadKeyIdentifierClauseCore(XmlReader reader)
        {
            if (reader.IsStartElement(CustomTokenSerializer.SecurityTokenReference, CustomTokenSerializer.WsSecurityNamespace))
            {
                string tokenType = reader.GetAttribute(CustomTokenSerializer.TokenType);
                if (tokenType == CustomTokenSerializer.EAISamlTokenType)
                    return true;
            }

            return base.CanReadKeyIdentifierClauseCore(reader);
        }

        protected override SecurityKeyIdentifierClause ReadKeyIdentifierClauseCore(XmlReader reader)
        {
            if (reader.IsStartElement(CustomTokenSerializer.SecurityTokenReference, CustomTokenSerializer.WsSecurityNamespace))
            {
                string tokenType = reader.GetAttribute(CustomTokenSerializer.TokenType, CustomTokenSerializer.Wsse11SecurityNamespace);
                if (tokenType == CustomTokenSerializer.EAISamlTokenType)
                {
                    reader.Read();
                    if (reader.IsStartElement(CustomTokenSerializer.KeyIdentifier, CustomTokenSerializer.WsSecurityNamespace))
                    {
                        string valueType = reader.GetAttribute(CustomTokenSerializer.ValueType);
                        if (valueType == CustomTokenSerializer.EAISamlValueType)
                        {
                            string assertionId = reader.ReadString();
                            reader.ReadEndElement(); // wsse:KeyIdentifier
                            reader.ReadEndElement(); // wsse:SecurityTokenReference
                            return new EAISamlAssertionKeyIdentifierClause(assertionId);
                        }
                    }

                    throw new InvalidOperationException("SAML 2.0 reference has a invalid format.");
                }
            }

            return base.ReadKeyIdentifierClauseCore(reader);
        }

        protected override bool CanWriteKeyIdentifierClauseCore(SecurityKeyIdentifierClause keyIdentifierClause)
        {
            return ((keyIdentifierClause is EAISamlAssertionKeyIdentifierClause) || base.CanWriteKeyIdentifierClauseCore(keyIdentifierClause));
        }

        protected override void WriteKeyIdentifierClauseCore(XmlWriter writer, SecurityKeyIdentifierClause keyIdentifierClause)
        {
            if (keyIdentifierClause is EAISamlAssertionKeyIdentifierClause)
            {
                writer.WriteStartElement(WsSecurityPrefix, SecurityTokenReference, WsSecurityNamespace);
                writer.WriteAttributeString(TokenType, Wsse11SecurityNamespace, EAISamlTokenType);
                writer.WriteStartElement(WsSecurityPrefix, KeyIdentifier, WsSecurityNamespace);
                writer.WriteAttributeString(ValueType, EAISamlValueType);
                writer.WriteString(((EAISamlAssertionKeyIdentifierClause)keyIdentifierClause).AssertionID);
                writer.WriteEndElement();
                writer.WriteEndElement();
            }
            else
            {
                base.WriteKeyIdentifierClauseCore(writer, keyIdentifierClause);
            }
        }
    }

    public class CustomServiceCredentials : ServiceCredentials
    {
        EAISamlSerializer samlSerializer;

        public CustomServiceCredentials()
            : base()
        {
        }

        CustomServiceCredentials(CustomServiceCredentials other)
            : base(other)
        {
        }

        public override SecurityTokenManager CreateSecurityTokenManager()
        {
            return new CustomServiceCredentialsTokenManager(this);
        }

        protected override ServiceCredentials CloneCore()
        {
            return new CustomServiceCredentials(this);
        }

        public EAISamlSerializer SamlSerializer
        {
            get
            {
                if (this.samlSerializer == null)
                    this.samlSerializer = new EAISamlSerializer();

                return this.samlSerializer;
            }
            set
            {
                this.samlSerializer = value;
            }
        }
    }

    class CustomServiceCredentialsTokenManager : ServiceCredentialsSecurityTokenManager
    {
        public CustomServiceCredentialsTokenManager(CustomServiceCredentials creds)
            : base(creds)
        {
        }

        public override SecurityTokenSerializer CreateSecurityTokenSerializer(SecurityTokenVersion version)
        {
            ReadOnlyCollection<string> specs = version.GetSecuritySpecifications();
            SecurityVersion secVersion;
            if (specs.Contains(CustomTokenSerializer.WsSecurityNamespace))
            {
                secVersion = SecurityVersion.WSSecurity10;
            }
            else
            {
                secVersion = SecurityVersion.WSSecurity11;
            }

            EAISamlSerializer samlSerializer = ((CustomServiceCredentials)this.ServiceCredentials).SamlSerializer;
            return new CustomTokenSerializer(secVersion, samlSerializer);
        }

        public override SecurityTokenAuthenticator CreateSecurityTokenAuthenticator(SecurityTokenRequirement tokenRequirement, out SecurityTokenResolver outOfBandTokenResolver)
        {
            if (tokenRequirement.TokenType == CustomTokenSerializer.EAISamlTokenType)
            {
                List<SecurityToken> tokens = new List<SecurityToken>();
                if (this.ServiceCredentials.ServiceCertificate.Certificate != null)
                    tokens.Add(new X509SecurityToken(this.ServiceCredentials.ServiceCertificate.Certificate));
                foreach (X509Certificate2 cert in this.ServiceCredentials.IssuedTokenAuthentication.KnownCertificates)
                {
                    tokens.Add(new X509SecurityToken(cert));
                }
                outOfBandTokenResolver = SecurityTokenResolver.CreateDefaultSecurityTokenResolver(tokens.AsReadOnly(), false);
                List<SecurityTokenAuthenticator> supportingTokenAuthenticator = new List<SecurityTokenAuthenticator>();
                supportingTokenAuthenticator.Add(new X509SecurityTokenAuthenticator(X509CertificateValidator.None));
                if (this.ServiceCredentials.IssuedTokenAuthentication.AllowUntrustedRsaIssuers)
                    supportingTokenAuthenticator.Add(new RsaSecurityTokenAuthenticator());
                return new EAISamlSecurityTokenAuthenticator(supportingTokenAuthenticator);
            }
            else
            {
                return base.CreateSecurityTokenAuthenticator(tokenRequirement, out outOfBandTokenResolver);
            }
        }
    }
}
